/**
 * 2017年5月28日
 */
package cn.edu.bjtu.io;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.util.Random;
import org.datavec.api.records.reader.impl.LineRecordReader;
import org.datavec.api.split.FileSplit;
import cn.edu.bjtu.core.LoggerSupport;

/**
 * @author Alex
 *
 */
public class FileDataSetSplit extends LoggerSupport{

	File in;
	File out;
	double ratio ;
	String trainSuffix;
	String testSuffix;
	int seed;
	protected FileDataSetSplit(Builder builder){
		this.in = builder.inputFile;
		this.out = builder.outputDir == null?in.getParentFile():builder.outputDir;
		this.ratio = builder.ratio;
		this.seed = builder.seed;
		this.trainSuffix = builder.trainSuffix;
		this.testSuffix = builder.testSuffix;
	}
	public void doSplit(){
		String outName = in.getName();
		Random rand = new Random(seed);
		try(LineRecordReader lrr = new LineRecordReader()){
				lrr.initialize(new FileSplit(in));
				try(BufferedWriter train = new BufferedWriter(new FileWriter(new File(out,outName+trainSuffix)))){
					try(BufferedWriter test = new BufferedWriter(new FileWriter(new File(out,outName+testSuffix)))){
						while(lrr.hasNext()){
							if(rand.nextDouble()>ratio){
								test.write(lrr.next().get(0).toString());
								test.newLine();
							}else{
								train.write(lrr.next().get(0).toString());
								train.newLine();
							}
						}
					}catch(Exception e){
						logException(e);
					}
				}catch(Exception e){
					logException(e);
				}
			
		}catch (Exception e) {
			logException(e);
		}
		
	}
	
	public static class Builder{
		File inputFile = null;
		File outputDir = null;
		double ratio = 0.8;
		int seed = 47;
		String trainSuffix = "_train";
		String testSuffix = "_test";
		public Builder() {}
		public Builder setInputFileOrDir(File file){
			inputFile = file;
			return this;
		}
		public Builder setRatio(double ratio){
			this.ratio = ratio;
			return this;
		}
		public Builder setOutputDir(File out){
			this.outputDir = out;
			if(out.isFile()){
				throw new IllegalArgumentException("OutPut should a DIRECTORY ");
			}
			return this;
		}
		public Builder setTrainSuffix(String suffix){
			this.trainSuffix = suffix;
			return this;
		}
		
		public Builder setTestSuffix(String suffix){
			this.testSuffix = suffix;
			return this;
		}
		
		public Builder setSeed(int seed){
			this.seed = seed;
			return this;
		}
		
		public FileDataSetSplit build(){
			return new FileDataSetSplit(this);
		}
	}
}
